package com.demon.framework.util.data;


import com.demon.framework.common.exception.QueryException;
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.parser.CCJSqlParserManager;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.select.*;
import net.sf.jsqlparser.util.TablesNamesFinder;
import org.springframework.jdbc.support.JdbcUtils;

import javax.sql.DataSource;
import java.io.StringReader;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.stream.Collectors;

/**
 * jdbc 操作 大量数据工具类
 *
 * @author Demon-HY
 * @date 2020-7-27
 */
@Slf4j
public class JdbcForBigDataUtil {

    /**
     * 根据sql解析出列名
     *
     * @param sql
     * @return
     */
    public static List<String> getColumnNames(String sql) {
        try {
            List<SelectItem> selectItems = new ArrayList<>();

            CCJSqlParserManager parser = new CCJSqlParserManager();
            Statement stmt = parser.parse(new StringReader(sql));
            if (stmt instanceof Select) {
                Select select = (Select) stmt;
                SelectBody selectBody = select.getSelectBody();
                //设置selectItems
                PlainSelect plainSelect = (PlainSelect) selectBody;

                for (SelectItem selectItem : plainSelect.getSelectItems()) {
                    //别名需要特殊处理
                    if (selectItem instanceof SelectExpressionItem) {
                        SelectExpressionItem selectExpressionItem = (SelectExpressionItem) selectItem;
                        if (selectExpressionItem.getAlias() != null) {
                            //直接使用别名
                            Column column = new Column(selectExpressionItem.getAlias().getName());
                            SelectExpressionItem expressionItem = new SelectExpressionItem(column);
                            selectItems.add(expressionItem);
                        } else if (selectExpressionItem.getExpression() instanceof Column) {
                            Column column = (Column) selectExpressionItem.getExpression();
                            SelectExpressionItem item;
                            if (column.getTable() != null) {
                                Column newColumn = new Column(column.getColumnName());
                                item = new SelectExpressionItem(newColumn);
                                selectItems.add(item);
                            } else {
                                selectItems.add(selectItem);
                            }
                        } else {
                            selectItems.add(selectItem);
                        }
                    } else if (selectItem instanceof AllTableColumns) {
                        selectItems.add(new AllColumns());
                    } else {
                        selectItems.add(selectItem);
                    }
                }
            }

            return selectItems.stream().map(selectItem -> selectItem.toString()).collect(Collectors.toList());
        } catch (Throwable e) {
            log.error("获取列名异常,sql:{}", sql, e);
            throw new QueryException("获取列名异常！", e);
        }

    }

    /**
     * 根据sql解析出表名
     *
     * @param sql sql的字段名为列名,支持重命名 ,重命名请不要使用引号 eg: select id 主键 from table
     * @return
     */
    public static List<String> getTableNames(String sql) {
        try {
            CCJSqlParserManager parser = new CCJSqlParserManager();
            List<String> list = new ArrayList<>();
            Statement stmt = parser.parse(new StringReader(sql));
            if (stmt instanceof Select) {
                Select selectStatement = (Select) stmt;
                TablesNamesFinder tablesNamesFinder = new TablesNamesFinder();
                List tableList = tablesNamesFinder.getTableList(selectStatement);
                for (Iterator iter = tableList.iterator(); iter.hasNext(); ) {
                    String tableName = iter.next().toString();
                    list.add(tableName);
                }
            }
            return list;
        } catch (Throwable e) {
            log.error("获取列名异常,sql:{}", sql, e);
            throw new QueryException("获取列名异常！", e);
        }

    }

    public static void main(String[] args) {
        List<String> columnNames = getColumnNames(" select id 主键,GOODS_ID as 不,KEY_PROTO as 要,SQL_NAME 引,VAL 号,CREATE_TIME from t_goods_prototype where goods_id=123");

        System.out.println(columnNames);
    }

    /**
     * 查询数据
     *
     * @param sql sql的字段名为列名,支持重命名 ,重命名请不要使用引号 eg: select id 主键 from table
     * @return
     * @See 大量数据时 流式 或 分批次多线程(TODO ) 查询DB
     */
    public static List<List<Object>> queryData(DataSource dataSource, String sql) {
        List<List<Object>> lists = Lists.newArrayList();

        try {
            List<String> columnNames = getColumnNames(sql);
            PreparedStatement preparedStatement = dataSource.getConnection().prepareStatement(sql);
            ResultSet rs = preparedStatement.executeQuery();
            while (rs.next()) {
                List<Object> data = Lists.newArrayList();

                for (String columnName : columnNames) {
                    data.add(rs.getObject(columnName));
                }

                lists.add(data);
            }

            return lists;
        } catch (Throwable e) {
            log.error("查询数据异常,sql:{}", sql, e);
            throw new QueryException("查询数据异常！", e);
        }

    }

    /**
     * 处理数据
     *
     * @param sql      sql的字段名为列名,支持重命名 ,重命名请不要使用引号 eg: select id 主键 from table
     * @param callable
     * @See 大量数据时 流式 或 分批次多线程(TODO ) 查询DB
     */
    public static void processData(DataSource dataSource, String sql, ProcessCallable callable) {
        Connection connection = null;
        PreparedStatement preparedStatement = null;
        ResultSet rs = null;
        try {
            connection = dataSource.getConnection();
            preparedStatement = connection.prepareStatement(sql,
                    ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY);
            //forward only read only也是mysql 驱动的默认值，所以不指定也是可以的 比如： PreparedStatement ps = connection.prepareStatement("select .. from ..");
            //也可以修改jdbc url通过defaultFetchSize参数来设置，这样默认所以的返回结果都是通过流方式读取.
            preparedStatement.setFetchSize(Integer.MIN_VALUE);
            rs = preparedStatement.executeQuery();

            List<String> columnNames = getColumnNames(sql);
            while (rs.next()) {
                List<Object> data = Lists.newArrayList();
                for (String columnName : columnNames) {
                    data.add(rs.getObject(columnName));
                }

                ArrayList<List<Object>> objects = new ArrayList<>();
                objects.add(data);

                callable.call(objects);
            }

        } catch (Throwable e) {
            log.error("处理数据异常,sql:{}", sql, e);
            throw new QueryException("处理数据异常！", e);
        } finally {
            JdbcUtils.closeResultSet(rs);
            JdbcUtils.closeStatement(preparedStatement);
            JdbcUtils.closeConnection(connection);
        }

    }

    /**
     * 数据处理的回调类
     */
    public interface ProcessCallable {
        void call(List<List<Object>> t) throws Exception;
    }
}
